type SessionUser = {
  role: string;
  [key: string]: unknown;
};

type AccessPolicy = {
  allow?: {
    roles?: string[];
  };
  deny?: {
    roles?: string[];
  };
};

type RegisteredToolWithAuth = {
  description?: string;
  inputSchema?: z.ZodObject<ZodRawShape>;
  callback: ToolCallback<ZodRawShape | undefined>;
  accessPolicy?: AccessPolicy;
};

// Just a simple extension to McpServer that adds support for access policies on tools
class McpServerWithAuth extends McpServer {
  protected override _registeredTools: {
    [name: string]: RegisteredToolWithAuth;
  } = {};
  checkPermissions(user?: SessionUser, policy?: AccessPolicy): boolean {
    if (!policy) {
      return true;
    }

    if (!user) {
      return false;
    }

    // Check deny rules first
    if (policy.deny) {
      // Check denied roles
      if (policy.deny.roles?.includes(user.role)) {
        return false;
      }
    }

    // Check allow rules
    if (policy.allow) {
      let isAllowed = false;

      // If no allow rules are specified, default to allowed
      if (!policy.allow.roles) {
        isAllowed = true;
      } else {
        // Check allowed roles
        if (policy.allow.roles?.includes(user.role)) {
          isAllowed = true;
        }
      }

      return isAllowed;
    }

    // If no rules specified, default to allowed
    return true;
  }

  override tool(
    name: string,
    cb: ToolCallback,
    accessPolicy?: AccessPolicy,
  ): void;
  override tool(
    name: string,
    description: string,
    cb: ToolCallback,
    accessPolicy?: AccessPolicy,
  ): void;
  override tool<Args extends ZodRawShape>(
    name: string,
    paramsSchema: Args,
    cb: ToolCallback<Args>,
    accessPolicy?: AccessPolicy,
  ): void;
  override tool<Args extends ZodRawShape>(
    name: string,
    description: string,
    paramsSchema: Args,
    cb: ToolCallback<Args>,
    accessPolicy?: AccessPolicy,
  ): void;
  override tool(name: string, ...rest: unknown[]): void {
    let description: string | undefined;
    let paramsSchema: ZodRawShape | undefined;
    let accessPolicy: AccessPolicy | undefined;
    let cb: ToolCallback<ZodRawShape | undefined>;

    // Parse arguments based on their types
    if (typeof rest[0] === "function") {
      // Case: tool(name, cb, accessPolicy?)
      cb = rest[0] as ToolCallback<ZodRawShape | undefined>;
      accessPolicy = rest[1] as AccessPolicy | undefined;
    } else if (typeof rest[0] === "string") {
      // Cases with description
      description = rest[0];
      if (typeof rest[1] === "function") {
        // Case: tool(name, description, cb, accessPolicy?)
        cb = rest[1] as ToolCallback<ZodRawShape | undefined>;
        accessPolicy = rest[2] as AccessPolicy | undefined;
      } else {
        // Case: tool(name, description, paramsSchema, cb, accessPolicy?)
        paramsSchema = rest[1] as ZodRawShape;
        cb = rest[2] as ToolCallback<ZodRawShape>;
        accessPolicy = rest[3] as AccessPolicy | undefined;
      }
    } else {
      // Case: tool(name, paramsSchema, cb, accessPolicy?)
      paramsSchema = rest[0] as ZodRawShape;
      cb = rest[1] as ToolCallback<ZodRawShape>;
      accessPolicy = rest[2] as AccessPolicy | undefined;
    }

    // Register with base class
    const args: unknown[] = [name];
    if (description) args.push(description);
    if (paramsSchema) args.push(paramsSchema);
    args.push(cb);

    // Set up request handlers if not already initialized
    if (!this._toolHandlersInitialized) {
      this.server.assertCanSetRequestHandler(
        CallToolRequestSchema.shape.method.value,
      );
      this.server.assertCanSetRequestHandler(
        ListToolsRequestSchema.shape.method.value,
      );
      this.server.registerCapabilities({ tools: {} });

      // Add ListToolsRequestSchema handler
      this.server.setRequestHandler(
        ListToolsRequestSchema,
        (request, extra): ListToolsResult => {
          const user = extra.user as SessionUser | undefined;

          // Filter tools based on permissions
          const accessibleTools = Object.entries(this._registeredTools)
            .filter(([_, tool]) =>
              this.checkPermissions(user, tool.accessPolicy),
            )
            .map(
              ([name, tool]): Tool => ({
                name,
                description: tool.description,
                inputSchema: tool.inputSchema
                  ? (zodToJsonSchema(tool.inputSchema, {
                      strictUnions: true,
                    }) as Tool["inputSchema"])
                  : { type: "object" },
              }),
            );

          return { tools: accessibleTools };
        },
      );

      this.server.setRequestHandler(
        CallToolRequestSchema,
        async (request: CallToolRequest, extra: RequestHandlerExtra) => {
          const tool = this._registeredTools[request.params.name];
          if (!tool) {
            throw new Error(`Tool ${request.params.name} not found`);
          }

          if (
            !this.checkPermissions(
              extra.user as SessionUser,
              tool.accessPolicy,
            )
          ) {
            throw new Error(`Access denied for tool: ${request.params.name}`);
          }

          if (tool.inputSchema) {
            const parseResult = await tool.inputSchema.safeParseAsync(
              request.params.arguments,
            );
            if (!parseResult.success) {
              throw new Error(
                `Invalid arguments for tool ${request.params.name}: ${parseResult.error.message}`,
              );
            }

            const args = parseResult.data;
            const cb = tool.callback as ToolCallback<ZodRawShape>;
            return await Promise.resolve(cb(args, extra));
          } else {
            const cb = tool.callback as ToolCallback<undefined>;
            return await Promise.resolve(cb(extra));
          }
        },
      );
      this._toolHandlersInitialized = true;
    }

    McpServer.prototype.tool.apply(
      this,
      args as Parameters<typeof McpServer.prototype.tool>,
    );
    this._registeredTools[name].accessPolicy = accessPolicy;
  }
}

const mcpServer = new McpServerWithAuth({
  name: "test server with auth",
  version: "1.0",
});
const client = new Client({
  name: "test client",
  version: "1.0",
});

mcpServer.tool("public-tool", async () => ({
  content: [
    {
      type: "text",
      text: "Public tool response",
    },
  ],
}));

mcpServer.tool(
  "protected-tool",
  async () => ({
    content: [
      {
        type: "text",
        text: "Protected tool response",
      },
    ],
  }),
  {
    allow: {
      roles: ["admin"],
    },
  },
);